import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset, random_split

# 数据加载和预处理
data = pd.read_csv("./data/心脏疾病数据集.csv").dropna()
data['bmi'] = data['weight'] / ((data['height'] / 100) ** 2)
data['age'] = data['age_year'].round().astype(int)

# 提取特征
features = ['age', 'gender', 'ap_hi', 'ap_lo', 'cholesterol', 'smoke', 'alco', 'active', 'bmi']
X = torch.tensor(data[features].values, dtype=torch.float32)
y = torch.tensor(data['cardio'].values, dtype=torch.float32).view(-1, 1)

# 构建 Dataset 和 DataLoader
dataset = TensorDataset(X, y)
train_size = int(0.8 * len(dataset))
train_ds, val_ds = random_split(dataset, [train_size, len(dataset) - train_size])
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=32)


# 构建神经网络模型
class HeartDiseaseModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(9, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.net(x)


model = HeartDiseaseModel()
criterion = nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 训练模型
for epoch in range(10):
    model.train()
    for xb, yb in train_loader:
        pred = model(xb)
        loss = criterion(pred, yb)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# 保存模型
torch.save(model.state_dict(), "./model/heart_model.pth")
print("模型已保存到./model/heart_model.pth")
